import numpy as np
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split


class MathLearn:
    """机器学习"""

    def __init__(self, algorithm, algorithm_param, test_size=0.7):
        self.algorithm = algorithm
        self.algorithm_param = algorithm_param
        self.test_size = test_size
        self.ai = None
        self.score = None

    def train(self, x, y):
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=self.test_size)
        self.ai = self.algorithm(**self.algorithm_param)
        self.ai.fit(x_train, y_train)

        y_predict = self.ai.predict(x_test)
        self.score = r2_score(y_test, y_predict)  # 回归评价，模型越好：r2→1；模型越差：r2→0

    def predict(self, x):
        return self.ai.predict(x)


if __name__ == '__main__':
    from sklearn.linear_model import LinearRegression, Ridge, Lasso, LogisticRegression
    from sklearn.mixture import GaussianMixture
    math_learn = MathLearn(LinearRegression, algorithm_param={})  # 线性回归
    x_sample, y_sample = np.arange(100).reshape((50, 2)), np.arange(100).reshape((50, 2))
    math_learn.train(x_sample, y_sample)
    print(math_learn.score)
    #
    # math_learn = MathLearn(Ridge, algorithm_param={})  #
    # x_sample, y_sample = np.arange(100).reshape((50, 2)), range(50)
    # math_learn.train(x_sample, y_sample)
    # print(math_learn.score)
    #
    # math_learn = MathLearn(Lasso, algorithm_param={'alpha': 0.1})  #
    # x_sample, y_sample = np.arange(100).reshape((50, 2)), range(50)
    # math_learn.train(x_sample, y_sample)
    # print(math_learn.score)

    # math_learn = MathLearn(LogisticRegression, {'solver': 'lbfgs', 'max_iter': 100})  #
    # x_sample, y_sample = np.arange(100).reshape((50, 2)), range(50)
    # math_learn.train(x_sample, y_sample)
    # print(math_learn.score)
